#include "wrapper_hypre.h"
#include <vector>

typedef HYPRE_Int idx_t;
typedef HYPRE_Real data_t;

const idx_t num_diag = 7;

// const idx_t dims_per_blk[] = {
//     40 , 15 , 25 , 
//     192, 20 , 40 , 
//     25 , 65 , 40 , 
//      4 , 25 , 40 , 
//      2 ,  9 ,  8 , 
//     40 , 25 ,  8 , 
//     40 , 25 , 25 , 
//     40 , 25 , 35 , 
//     15 ,  5 , 40 , 
//      5 ,  3 , 40 , 
//      3 , 15 , 40 , 
//      2 ,  9 , 40 , 
//     40 , 25 , 40 , 
// };
// const idx_t offd_nnz [] = {
//       0,    600,   1000,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0,
//     600,      0,   2600,    160,      0,    320,   1000,   1400,      0,      0,      0,      0,   1600,
//    1000,   2600,      0,      0,      0,      0,   1000,      0,      0,      0,      0,      0,      0,
//       0,    160,      0,      0,      0,   1000,      0,   1000,      0,      0,      0,      0,      0,
//       0,      0,      0,      0,      0,     72,      0,      0,      0,      0,      0,     18,      0,
//       0,    320,      0,   1000,     72,      0,      0,      0,      0,      0,      0,      0,   1000,
//       0,   1000,   1000,      0,      0,      0,      0,   1000,      0,      0,      0,      0,      0,
//       0,   1400,      0,   1000,      0,      0,   1000,      0,      0,      0,      0,      0,      0,
//       0,      0,      0,      0,      0,      0,      0,      0,      0,    200,    600,      0,      0,
//       0,      0,      0,      0,      0,      0,      0,      0,    200,      0,    120,      0,    200,
//       0,      0,      0,      0,      0,      0,      0,      0,    600,    120,      0,      0,    600,
//       0,      0,      0,      0,     18,      0,      0,      0,      0,      0,      0,      0,    360,
//       0,   1600,      0,      0,      0,   1000,      0,      0,      0,    200,    600,    360,      0,
// };

// const idx_t dims_per_blk[] = {
//     //  80,  25,   50,
//     // 374,  35,   80,
//     //  50, 130,   80,
//     //   8,  50,   80,
//     //   4,  19,   16,
//     //  80,  50,   16,
//     //  80,  50,   50,
//     //  80,  50,   70,
//     //  29,  11,   75,
//     //  11,   5,   75,
//     //   5,  29,   75,
//     //   4,  19,   75,
//     //  75,  50,   80,
//  25,  50, 80,
//  35, 374, 80,
// 130,  50, 80,
//   8,  50, 80,
//  16,   4, 19,
//  50,  16, 80,
//  50,  80, 50,
//  70,  80, 50,
//  75,  29, 11,
//  75,   5, 11,
//  75,  29, 5,
//  75,   4, 19,
//  75,  50, 80,
// };
// const idx_t offd_nnz [] = {
//       0,    2000,    4000,       0,       0,       0,       0,       0,       0,       0,       0,       0,       0, 
//    2000,       0,   10400,     640,       0,    1280,    4000,    5600,       0,       0,       0,       0,    6000, 
//    4000,   10400,       0,       0,       0,       0,    4000,       0,       0,       0,       0,       0,       0, 
//       0,     640,       0,       0,       0,    4000,       0,    4000,       0,       0,       0,       0,       0, 
//       0,       0,       0,       0,       0,     304,       0,       0,       0,       0,       0,      76,       0, 
//       0,    1280,       0,    4000,     304,       0,       0,       0,       0,       0,       0,       0,    4000, 
//       0,    4000,    4000,       0,       0,       0,       0,    4000,       0,       0,       0,       0,       0, 
//       0,    5600,       0,    4000,       0,       0,    4000,       0,       0,       0,       0,       0,       0, 
//       0,       0,       0,       0,       0,       0,       0,       0,       0,     825,    2175,       0,       0, 
//       0,       0,       0,       0,       0,       0,       0,       0,     825,       0,     375,       0,     825, 
//       0,       0,       0,       0,       0,       0,       0,       0,    2175,     375,       0,       0,    2175, 
//       0,       0,       0,       0,      76,       0,       0,       0,       0,       0,       0,       0,    1425, 
//       0,    6000,       0,       0,       0,    4000,       0,       0,       0,     825,    2175,    1425,       0, 
// };

const idx_t dims_per_blk[] = {
230, 60, 108,
230,60, 108,
104, 24, 5,
39,	15,	80,
15,	5,	80,
5,	39,	80,
460, 150, 108,
};
const idx_t offd_nnz [] = {
      0,    6480,       0,       0,       0,       0,   24840, // 31320
   6480,       0,    2496,       0,    1200,    3120,   24840, // 38136
      0,    2496,       0,       0,       0,       0,       0, //  2496
      0,       0,       0,       0,    1200,    3120,       0, //  4320
      0,    1200,       0,    1200,       0,     400,       0, //  2800
      0,    3120,       0,    3120,     400,       0,       0, //  6640
  24840,   24840,       0,       0,       0,       0,       0, // 49680
};

int main(int argc, char* argv[])
{
    int my_pid, num_procs;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);

    idx_t arg_cnt = 1;
    const std::string pathname(argv[arg_cnt++]);
    const std::string assign_file(argv[arg_cnt++]);
    const std::string solver_name = std::string(argv[arg_cnt++]),
                        prec_name = std::string(argv[arg_cnt++]);
    const idx_t glb_nparts = sizeof(dims_per_blk) / sizeof(idx_t) / 3;
    idx_t blk_offsets[glb_nparts + 1];
    blk_offsets[0] = 0;
    for (idx_t b = 0; b < glb_nparts; b++) {
        blk_offsets[b + 1] = blk_offsets[b] 
            + dims_per_blk[b*3  ] * dims_per_blk[b*3+1] * dims_per_blk[b*3+2];
    }
    const idx_t tot_nrows = blk_offsets[glb_nparts];// assert(tot_nrows == 351864);

    idx_t obj_type = HYPRE_SSTRUCT;
    if (prec_name == "AMG") obj_type = HYPRE_PARCSR;
    if (my_pid == 0) {
        printf("object type : %d\n", obj_type);
    }

    std::vector<PartDescriptor<idx_t> > my_parts, glb_parts;
    {// 添加本进程所拥有的块和盒子
        FILE * fp = fopen(assign_file.c_str(), "r");
        idx_t tmp;
        fscanf(fp, "%d", &tmp); assert(tmp == glb_nparts);
        for (idx_t b = 0; b < glb_nparts; b++) {
            idx_t blk_id, blk_exts[3], blk_nprocs;
            fscanf(fp, "%d (%d,%d,%d) %d", &blk_id, &blk_exts[0], &blk_exts[1], &blk_exts[2], &blk_nprocs);
            assert(b == blk_id);
            {// 记录全局块信息
                PartDescriptor<idx_t> descr;
                descr.id = b;
                descr.nelems = 1;
                const idx_t * dims = dims_per_blk + b * 3;
                for (idx_t d = 0; d < 3; d++) {
                    descr.ilower[d] = 0;
                    // descr.iupper[d] = dims[2 - d] - 1;// 注意这里要减1
                    descr.iupper[d] = dims[d] - 1;
                    descr.nelems *= (descr.iupper[d] - descr.ilower[d] + 1);
                }
                glb_parts.emplace_back(descr);
            }
            // 检查自己是否有参与该块的计算
            for (idx_t p = 0; p < blk_nprocs; p++) {// 读取这么多次
                idx_t pid, begs[3], ends[3];
                fscanf(fp, "%d (%d,%d,%d) (%d,%d,%d)", &pid, &begs[0], &begs[1], &begs[2], &ends[0], &ends[1], &ends[2]);
                if (pid == my_pid) {// 是本进程负责的盒子
                    PartDescriptor<idx_t> descr;
                    descr.id = b;
                    descr.nelems = 1;
                    for (idx_t d = 0; d < 3; d++) {
                        descr.ilower[d] = begs[d];
                        descr.iupper[d] = ends[d] - 1;
                        descr.nelems *= (descr.iupper[d] - descr.ilower[d] + 1);
                    }
                    my_parts.emplace_back(descr);
                }
            }
        }
        fclose(fp); 
        // 确定本进程在每块中都只负责一个盒子
        idx_t cnts[glb_nparts];
        for (idx_t g = 0; g < glb_nparts; g++) cnts[g] = 0;
        for (idx_t g = 0; g < (idx_t) my_parts.size(); g++) {
            cnts[my_parts[g].id] ++;
        }
        for (idx_t g = 0; g < glb_nparts; g++) assert(cnts[g] <= 1);
    }
    const idx_t my_nparts = my_parts.size();
#ifdef _zyDEBUG
    // for (int p = 0; p < num_procs; p++) {
    //     if (my_pid == p) {
    //         printf("Proc %d own %d parts:\n", my_pid, my_nparts);
    //         for (idx_t mp = 0; mp < my_nparts; mp++) {
    //             const PartDescriptor<idx_t> & descr = my_parts[mp];
    //             printf("\t Part%3d (%3d,%3d,%3d)~(%3d,%3d,%3d)\n", descr.id,
    //                 descr.ilower[0], descr.ilower[1], descr.ilower[2],
    //                 descr.iupper[0], descr.iupper[1], descr.iupper[2]  );
    //         }
    //         fflush(stdout);
    //     }
    //     MPI_Barrier(MPI_COMM_WORLD);
    // }
    // MPI_Barrier(MPI_COMM_WORLD);
#endif

    HYPRE_Init();

    HYPRE_SStructGrid ssgrid;
    HYPRE_SStructStencil stencil;
    HYPRE_SStructGraph ssgraph;
    HYPRE_SStructMatrix   A;
    HYPRE_SStructVector   b, x, y, x0;
    HYPRE_ParCSRMatrix par_A;
    HYPRE_ParVector par_b, par_x, par_y;
    TEST_RECORD records[TEST_CNT];
    TEST_CONFIG config;
    config.case_name = "kvlcc2";
    config.rtol = 1.0e-3;

    {
        // 创建网格
        HYPRE_SStructGridCreate(MPI_COMM_WORLD, 3, glb_nparts, &ssgrid);
        for (idx_t mb = 0; mb < my_nparts; mb++) {
            HYPRE_SStructGridSetExtents(ssgrid, my_parts[mb].id, my_parts[mb].ilower, my_parts[mb].iupper);
        }
        HYPRE_SStructVariable vtypes[1] = { HYPRE_SSTRUCT_VARIABLE_CELL };
        for (idx_t part = 0; part < glb_nparts; part++) {// 每个part都要执行定义变量的操作
            HYPRE_SStructGridSetVariables(ssgrid, part, sizeof(vtypes)/sizeof(HYPRE_SStructVariable), vtypes);
        }
        HYPRE_SStructGridAssemble(ssgrid);// a collective call finalizing the grid assembly.
        // 创建stencil
        HYPRE_SStructStencilCreate(3, num_diag, &stencil);
        idx_t offsets[7][3] = {
            {-1,0,0}, {0,-1,0}, {0,0,-1}, {0,0,0}, {0,0,1}, {0,1,0}, {1,0,0},// 正确
            // {-1,0,0}, {0,0,-1}, {0,-1,0}, {0,0,0}, {0,1,0}, {0,0,1}, {1,0,0},// 差一些
            // {0,0,-1}, {-1,0,0}, {0,-1,0}, {0,0,0}, {0,1,0}, {1,0,0}, {0,0,1},// 差很多
            // {0,0,-1}, {0,-1,0}, {-1,0,0}, {0,0,0}, {1,0,0}, {0,1,0}, {0,0,1},// 差很多
            // {0,-1,0}, {0,0,-1}, {-1,0,0}, {0,0,0}, {1,0,0}, {0,0,1}, {0,1,0},// 差很多
            // {0,-1,0}, {-1,0,0}, {0,0,-1}, {0,0,0}, {0,0,1}, {1,0,0}, {0,1,0},// 差很多

            // {0,0,1}, {0,1,0}, {1,0,0}, {0,0,0}, {-1,0,0}, {0,-1,0}, {0,0,-1}, 
        };
        for (idx_t e = 0; e < num_diag; e++) {
            HYPRE_SStructStencilSetEntry(stencil, e, offsets[e], 0);// 只有0号变量一种
        }
        // 创建图
        HYPRE_SStructGraphCreate(MPI_COMM_WORLD, ssgrid, &ssgraph);
        HYPRE_SStructGraphSetObjectType(ssgraph, obj_type);
        for (idx_t part = 0; part < glb_nparts; part++) {// 每个part都要执行定义模板的操作
            HYPRE_SStructGraphSetStencil(ssgraph, part, 0, stencil);
        }
        // 添加上非stencil的连接关系
        for (idx_t mp = 0; mp < my_nparts; mp++) {
            const PartDescriptor<idx_t> & loc_info = glb_parts[my_parts[mp].id];
            // TODO: 在目录下寻找名为"B%d_Offd*"的文件，并从文件大小确定有多少个非零元
            FILE * fp = nullptr;
            for (idx_t np = 0; np < glb_nparts; np++) {// 遍历每一个可能的邻居块
                const idx_t nnz = offd_nnz[loc_info.id * glb_nparts + np];
                if (nnz == 0) continue;
                // 读入该非对角块的数据
                const PartDescriptor<idx_t> & ngb_info = glb_parts[np];

                std::vector<COO<idx_t, data_t> > buf(nnz);
                fp = fopen((pathname + "/B"+std::to_string(loc_info.id)+"_Offd"+std::to_string(ngb_info.id)).c_str(), "rb");
                size_t ret = fread(buf.data(), sizeof(COO<idx_t, data_t>), nnz, fp);
                assert(ret == (size_t) nnz);
                fclose(fp);

                idx_t loc_dims[3], ngb_dims[3];
                for (idx_t d = 0; d < 3; d++) {
                    loc_dims[d] = loc_info.iupper[d] - loc_info.ilower[d] + 1;
                    ngb_dims[d] = ngb_info.iupper[d] - ngb_info.ilower[d] + 1;
                }
                // 为该非对角块设置连接关系
                for (idx_t k = 0; k < nnz; k++) {
                    const idx_t loc_1D = buf[k].r - blk_offsets[loc_info.id],// 该非零元的行序号在本块内的一维索引
                                ngb_1D = buf[k].c - blk_offsets[ngb_info.id];// 该非零元的列序号在邻居块内的一维索引
                    assert(loc_1D >= 0 && loc_1D < loc_info.nelems);
                    assert(ngb_1D >= 0 && ngb_1D < ngb_info.nelems);
                    idx_t index[3], to_id[3];
                    index[0] = (loc_1D                                       ) / (loc_dims[1] * loc_dims[2]);
                    index[1] = (loc_1D - index[0] * loc_dims[1] * loc_dims[2]) / (loc_dims[2]              );
                    index[2] = (loc_1D - index[0] * loc_dims[1] * loc_dims[2] - index[1] * loc_dims[2]);
                    if (my_parts[mp].ilower[0] <= index[0] && index[0] <= my_parts[mp].iupper[0] &&
                        my_parts[mp].ilower[1] <= index[1] && index[1] <= my_parts[mp].iupper[1] &&
                        my_parts[mp].ilower[2] <= index[2] && index[2] <= my_parts[mp].iupper[2] )// 属于本进程负责的范围内
                    {
                        to_id[0] = (ngb_1D                                       ) / (ngb_dims[1] * ngb_dims[2]);
                        to_id[1] = (ngb_1D - to_id[0] * ngb_dims[1] * ngb_dims[2]) / (ngb_dims[2]              );
                        to_id[2] = (ngb_1D - to_id[0] * ngb_dims[1] * ngb_dims[2] - to_id[1] * ngb_dims[2]);

                        HYPRE_SStructGraphAddEntries(ssgraph, loc_info.id, index, 0, ngb_info.id, to_id, 0);
                    }
                }// nnz loop
            }// ngb part loop
        }// my part loop
        HYPRE_SStructGraphAssemble(ssgraph);

        // 建立向量
        HYPRE_SStructVectorCreate(MPI_COMM_WORLD, ssgrid, &b);// Create an empty vector object
        HYPRE_SStructVectorCreate(MPI_COMM_WORLD, ssgrid, &x);
        HYPRE_SStructVectorCreate(MPI_COMM_WORLD, ssgrid, &y);
        HYPRE_SStructVectorCreate(MPI_COMM_WORLD, ssgrid, &x0);
        HYPRE_SStructVectorSetObjectType(b, obj_type);// Set the object type for the vectors to be the same as was already set for the matrix
        HYPRE_SStructVectorSetObjectType(x, obj_type);
        HYPRE_SStructVectorSetObjectType(y, obj_type);
        HYPRE_SStructVectorSetObjectType(x0, obj_type);
        HYPRE_SStructVectorInitialize(b);// Indicate that the vector coefficients are ready to be set
        HYPRE_SStructVectorInitialize(x);
        HYPRE_SStructVectorInitialize(y);
        HYPRE_SStructVectorInitialize(x0);
        {// 填充向量数据
            FILE * fp = nullptr; size_t ret;
            for (idx_t mp = 0; mp < my_nparts; mp++) {
                PartDescriptor<idx_t> & loc_part = my_parts[mp];
                PartDescriptor<idx_t> & loc_info = glb_parts[loc_part.id];

                const idx_t tot_nelems = blk_offsets[loc_info.id + 1] - blk_offsets[loc_info.id]; assert(tot_nelems == loc_info.nelems);
                const idx_t  my_nelems = loc_part.nelems;

                data_t * vec_buf = new data_t [tot_nelems];// [dim0][dim1][dim2]
                data_t * set_buf = new data_t [ my_nelems];// [dim2][dim1][dim0]

                fp = fopen((pathname + "/b.B" + std::to_string(loc_info.id)).c_str(), "rb");
                ret = fread(vec_buf, sizeof(data_t), tot_nelems, fp); assert(ret == (size_t) tot_nelems);
                fclose(fp);

                idx_t tot_dims[3], my_dims[3];
                for (idx_t d = 0; d < 3; d++) {
                    tot_dims[d] = loc_info.iupper[d] - loc_info.ilower[d] + 1;
                     my_dims[d] = loc_part.iupper[d] - loc_part.ilower[d] + 1;
                }

                #pragma omp parallel for collapse(2) schedule(static)
                for (idx_t i0 = 0; i0 < my_dims[0]; i0++)
                for (idx_t i1 = 0; i1 < my_dims[1]; i1++)
                for (idx_t i2 = 0; i2 < my_dims[2]; i2++)
                    set_buf[((i2 *  my_dims[1]) + i1) *  my_dims[0] + i0] = 
                    vec_buf[(((i0 + loc_part.ilower[0]) * tot_dims[1]) + i1 + loc_part.ilower[1]) * tot_dims[2] + i2 + loc_part.ilower[2]];

                HYPRE_SStructVectorSetBoxValues(b, loc_part.id, loc_part.ilower, loc_part.iupper, 0, set_buf);

                fp = fopen((pathname + "/x0.B" + std::to_string(loc_part.id)).c_str(), "rb");
                ret = fread(vec_buf, sizeof(data_t), tot_nelems, fp); assert(ret == (size_t) tot_nelems);
                fclose(fp);
                
                #pragma omp parallel for collapse(2) schedule(static)
                for (idx_t i0 = 0; i0 < my_dims[0]; i0++)
                for (idx_t i1 = 0; i1 < my_dims[1]; i1++)
                for (idx_t i2 = 0; i2 < my_dims[2]; i2++)
                    set_buf[((i2 *  my_dims[1]) + i1) *  my_dims[0] + i0] = 
                    vec_buf[(((i0 + loc_part.ilower[0]) * tot_dims[1]) + i1 + loc_part.ilower[1]) * tot_dims[2] + i2 + loc_part.ilower[2]];

                HYPRE_SStructVectorSetBoxValues(x, loc_part.id, loc_part.ilower, loc_part.iupper, 0, set_buf);

                #pragma omp parallel for schedule(static)
                for (idx_t i = 0; i < my_nelems; i++) set_buf[i] = 0.0;
                HYPRE_SStructVectorSetBoxValues(y, loc_part.id, loc_part.ilower, loc_part.iupper, 0, set_buf);

                delete [] vec_buf; vec_buf = nullptr;
                delete [] set_buf; set_buf = nullptr;
            }// my part loop
        }// fill in vector data
        HYPRE_SStructVectorAssemble(b);
        HYPRE_SStructVectorAssemble(x);
        HYPRE_SStructVectorAssemble(y);// 对y也要创建
        HYPRE_SStructVectorAssemble(x0);
        HYPRE_SStructVectorCopy(x, x0);
        
        HYPRE_SStructMatrixCreate(MPI_COMM_WORLD, ssgraph, &A);// Create an empty matrix object
        HYPRE_SStructMatrixSetObjectType(A, obj_type);
        HYPRE_SStructMatrixInitialize(A);// Indicate that the matrix coefficients are ready to be set 
        idx_t stencil_indices[num_diag];
        for (idx_t j = 0; j < num_diag; j++)
            stencil_indices[j] = j;
        {// 填充矩阵数据
            FILE * fp = nullptr; size_t ret;
            for (idx_t mp = 0; mp < my_nparts; mp++) {
                PartDescriptor<idx_t> & loc_part = my_parts[mp];
                PartDescriptor<idx_t> & loc_info = glb_parts[loc_part.id];
                // 先设置对角块的数值
                const idx_t tot_nelems = blk_offsets[loc_info.id + 1] - blk_offsets[loc_info.id]; assert(tot_nelems == loc_info.nelems);
                const idx_t  my_nelems = loc_part.nelems;

                idx_t tot_dims[3], my_dims[3];
                for (idx_t d = 0; d < 3; d++) {
                    tot_dims[d] = loc_info.iupper[d] - loc_info.ilower[d] + 1;
                     my_dims[d] = loc_part.iupper[d] - loc_part.ilower[d] + 1;
                }

                data_t * mat_buf = new data_t [tot_nelems * num_diag];
                data_t * set_buf = new data_t [ my_nelems * num_diag];

                fp = fopen((pathname + "/B" + std::to_string(loc_info.id) + "_Diag").c_str(), "rb");
                ret = fread(mat_buf, sizeof(data_t), tot_nelems * num_diag, fp); assert(ret == (size_t)(tot_nelems * num_diag));
                fclose(fp);
                #pragma omp parallel for collapse(2) schedule(static)
                for (idx_t i0 = 0; i0 < my_dims[0]; i0++)
                for (idx_t i1 = 0; i1 < my_dims[1]; i1++)
                for (idx_t i2 = 0; i2 < my_dims[2]; i2++) {
                    data_t       * dst_ptr = set_buf + (((i2 *  my_dims[1]) + i1) *  my_dims[0] + i0) * num_diag;
                    const data_t * src_ptr = mat_buf + ((((i0 + loc_part.ilower[0]) * tot_dims[1]) + i1 + loc_part.ilower[1]) * tot_dims[2] + i2 + loc_part.ilower[2]) * num_diag;
                    for (idx_t d = 0; d < num_diag; d++)
                        dst_ptr[d] = src_ptr[d];
                    
                    // if (dims[0] == 40 && dims[1] == 15 && dims[2] == 25) {
                    //     for (idx_t d = 0; d < num_diag; d++)
                    //         printf("%.5e ", src_ptr[d]);
                    //     printf("\n");
                    // }
                }

                HYPRE_SStructMatrixSetBoxValues(A, loc_part.id, loc_part.ilower, loc_part.iupper, 0,
                    num_diag, stencil_indices, set_buf);
                delete [] mat_buf; mat_buf = nullptr;
                delete [] set_buf; set_buf = nullptr;
                // 再逐个设置非对角块的数值
                idx_t * nentries = new idx_t [my_nelems];// 记录已经为该位置设置了多少个非零元
                #pragma omp parallel for schedule(static)
                for (idx_t i = 0; i < my_nelems; i++) nentries[i] = num_diag;

                for (idx_t ngb = 0; ngb < glb_nparts; ngb++) {// 遍历每一个可能的邻居块
                    const idx_t nnz = offd_nnz[loc_info.id * glb_nparts + ngb];
                    if (nnz == 0) continue;
                    // 读入该非对角块的数据
                    PartDescriptor<idx_t> & ngb_info = glb_parts[ngb];

                    std::vector<COO<idx_t, data_t> > buf(nnz);
                    fp = fopen((pathname + "/B"+std::to_string(loc_info.id)+"_Offd"+std::to_string(ngb_info.id)).c_str(), "rb");
                    size_t ret = fread(buf.data(), sizeof(COO<idx_t, data_t>), nnz, fp);
                    assert(ret == (size_t) nnz);
                    fclose(fp);

                    idx_t loc_dims[3], ngb_dims[3];
                    for (idx_t d = 0; d < 3; d++) {
                        loc_dims[d] = loc_info.iupper[d] - loc_info.ilower[d] + 1;
                        ngb_dims[d] = ngb_info.iupper[d] - ngb_info.ilower[d] + 1;
                    }
                    // 为该非对角块设置连接关系
                    for (idx_t k = 0; k < nnz; k++) {
                        const idx_t loc_1D = buf[k].r - blk_offsets[loc_part.id];// 该非零元的行序号在本块内的一维索引
                        idx_t index[3];
                        index[0] = (loc_1D                                       ) / (loc_dims[1] * loc_dims[2]);
                        index[1] = (loc_1D - index[0] * loc_dims[1] * loc_dims[2]) / (loc_dims[2]              );
                        index[2] = (loc_1D - index[0] * loc_dims[1] * loc_dims[2] - index[1] * loc_dims[2]);
                        if (my_parts[mp].ilower[0] <= index[0] && index[0] <= my_parts[mp].iupper[0] &&
                            my_parts[mp].ilower[1] <= index[1] && index[1] <= my_parts[mp].iupper[1] &&
                            my_parts[mp].ilower[2] <= index[2] && index[2] <= my_parts[mp].iupper[2] )// 属于本进程负责的范围内
                        {
                            const idx_t my_1D = ((index[0] - loc_part.ilower[0]) * my_dims[1] + index[1] - loc_part.ilower[1]) * my_dims[2] + index[2] - loc_part.ilower[2];
                            HYPRE_SStructMatrixSetValues(A, loc_part.id, index, 0, 1, &nentries[my_1D], &buf[k].v);
                            nentries[my_1D] ++;// 计数增加
                        }
                    }// nnz loop
                }// ngb loop
                delete [] nentries; nentries = nullptr;
            }// my part loop
        }// fill in matrix data
        HYPRE_SStructMatrixAssemble(A);// a collective call finalizing the matrix assembly.
    
            for (idx_t test = 0; test < TEST_CNT; test++) {
            HYPRE_SStructVectorCopy(x0, x);
            
            data_t b_dot, x_dot, Ab_dot, Ax_dot;
            idx_t ret;
            // 做spmv检验一下数据传对没有
            if (obj_type == HYPRE_PARCSR) {
                HYPRE_SStructMatrixGetObject(A, (void **) &par_A);
                HYPRE_SStructVectorGetObject(b, (void **) &par_b);
                HYPRE_SStructVectorGetObject(x, (void **) &par_x);
                HYPRE_SStructVectorGetObject(y, (void **) &par_y);
                ret = HYPRE_ParVectorInnerProd(par_b, par_b, &b_dot);
                ret = HYPRE_ParVectorInnerProd(par_x, par_x, &x_dot);
                HYPRE_ParCSRMatrixMatvec(1.0, par_A, par_b, 0.0, par_y);
                ret = HYPRE_ParVectorInnerProd(par_y, par_y, &Ab_dot);
                HYPRE_ParCSRMatrixMatvec(1.0, par_A, par_x, 0.0, par_y);
                ret = HYPRE_ParVectorInnerProd(par_y, par_y, &Ax_dot);
            } else { assert(obj_type == HYPRE_SSTRUCT);
                ret = HYPRE_SStructInnerProd(b, b, &b_dot);
                ret = HYPRE_SStructInnerProd(x, x, &x_dot);
                HYPRE_SStructMatrixMatvec(1.0, A, b, 0.0, y);
                ret = HYPRE_SStructInnerProd(y, y, &Ab_dot);
                HYPRE_SStructMatrixMatvec(1.0, A, x, 0.0, y);
                ret = HYPRE_SStructInnerProd(y, y, &Ax_dot);
            }

            /*
            {
                data_t* vec_buf = new data_t [local_m * local_m * local_m];
                data_t* mat_buf = new data_t [local_m * local_m * local_m * num_diag];
                HYPRE_SStructVectorGetBoxValues(y, my_part, my_ilower, my_iupper, 0, vec_buf);
                HYPRE_SStructMatrixGetBoxValues(A, my_part, my_ilower, my_iupper, 0, num_diag, stencil_indices, mat_buf);
                for (int p = 0; p < num_procs; p++) {
                    if (p == my_pid) {
                        for (idx_t i2 = my_ilower[2]; i2 <= my_iupper[2]; i2++)
                        for (idx_t i1 = my_ilower[1]; i1 <= my_iupper[1]; i1++)
                        for (idx_t i0 = my_ilower[0]; i0 <= my_iupper[0]; i0++) {
                            const idx_t elem_id = ((i2 - my_ilower[2]) * (my_iupper[1] - my_ilower[1] + 1)
                                                    +   i1 - my_ilower[1] )* (my_iupper[0] - my_ilower[0] + 1)
                                                    +   i0 - my_ilower[0]  ;
                            const data_t tmp = vec_buf[elem_id];
                            const data_t* mat_vals = mat_buf + elem_id * num_diag;
                            // if (tmp != 0.0) {
                                printf("Part %d (%d,%d,%d): %.2f   A %.2f %.2f %.2f %.2f %.2f %.2f %.2f\n",
                                    my_part, i0, i1, i2, tmp,
                                    mat_vals[0], mat_vals[1], mat_vals[2], mat_vals[3], mat_vals[4], mat_vals[5], mat_vals[6]);
                            // }
                            // assert(tmp <= 3.0);
                        }
                        fflush(stdout);
                    }
                    MPI_Barrier(MPI_COMM_WORLD);
                }
            }*/
            

            if (my_pid == 0) {
                printf("(  b,   b) = %.20e\n",  (double)b_dot);
                printf("(  x,   x) = %.20e\n",  (double)x_dot);
                printf("(A*b, A*b) = %.20e\n", (double)Ab_dot);
                printf("(A*x, A*x) = %.20e\n", (double)Ax_dot);
            }

            // 计算初始残差以便确定收敛阈值
            data_t r_nrm2 = 0.0, b_nrm2 = 0.0;
            // check_residual(A, x, b, y, r_nrm2, b_nrm2);
            // config.atol = r_nrm2 * 1.0e-6;

            buildup_solver(solver_name, prec_name, config);

            idx_t num_iterations = 0;
            data_t t_setup = 0.0, t_solve = 0.0, final_res_norm = 0.0;

            if (obj_type == HYPRE_SSTRUCT)
                setup_and_solve(solver_name, prec_name, A, b, x, t_setup, t_solve, final_res_norm, num_iterations);
            else if (obj_type == HYPRE_PARCSR) {
                setup_and_solve(solver_name, prec_name, par_A, par_b, par_x, t_setup, t_solve, final_res_norm, num_iterations);
                HYPRE_SStructVectorGather(b);
                HYPRE_SStructVectorGather(x);
            }

            // 计算真实残差
            check_residual(prec_name, A, x, b, y, r_nrm2, b_nrm2);
            if (my_pid == 0) {
                printf("\033[1;35mtrue ||r|| = %20.16e ||r||/||b||= %20.16e\033[0m\n", r_nrm2, r_nrm2/b_nrm2);
                printf("Iterations = %d\n", num_iterations);
                printf("Time cost %.5f %.5f %.5f %d\n", (double)t_setup, (double)t_solve, (double)(t_setup + t_solve), num_iterations);
                printf("Final Relative Residual Norm = %e\n", (double)final_res_norm);
                printf("\n");
            }

            records[test].iter = num_iterations;
            records[test].setup = t_setup;
            records[test].solve = t_solve;

            destroy_solver(solver_name, prec_name);
        }// test loop

        if (my_pid == 0) {
            std::sort(records, records + TEST_CNT);
            const TEST_RECORD & best = records[0], & worst = records[TEST_CNT - 1];
            TEST_RECORD avg;
            for (idx_t i = 0; i < TEST_CNT; i++) {
                avg.setup += records[i].setup;
                avg.solve += records[i].solve;
                avg.prec  += records[i].prec;
                avg.iter  += records[i].iter;
            }
            avg.setup /= TEST_CNT;
            avg.solve /= TEST_CNT;
            avg.prec  /= TEST_CNT;
            avg.iter  /= TEST_CNT;
            data_t min_tot = best .setup + best .solve;
            data_t max_tot = worst.setup + worst.solve;
            data_t avg_tot = avg  .setup + avg  .solve;
            printf("\n  Summary of %d tests:\n", TEST_CNT);
            printf("     Setup time  Solve time  #Iter  Total time\n");
            printf("Min    %.6f    %.6f  %5d  %.6f\n", best .setup, best .solve, best .iter, min_tot);
            printf("Avg    %.6f    %.6f  %5d  %.6f\n", avg  .setup, avg  .solve, avg  .iter, avg_tot);
            printf("Max    %.6f    %.6f  %5d  %.6f\n", worst.setup, worst.solve, worst.iter, max_tot);    
        }

        HYPRE_SStructMatrixDestroy(A);
        HYPRE_SStructVectorDestroy(b); HYPRE_SStructVectorDestroy(x); HYPRE_SStructVectorDestroy(y);
        HYPRE_SStructGraphDestroy(ssgraph);
        HYPRE_SStructStencilDestroy(stencil);
        HYPRE_SStructGridDestroy(ssgrid);
    }

    HYPRE_Finalize();
    MPI_Finalize();
    return 0;
}